from util.pic_video import *
from util.LabelPoints2ImagebyMayaParam import *
import os
import time
import traceback
import sys



# pts = np.concatenate([pts, np.zeros((pts.shape[0],1),dtype=np.float32)],axis=1)


# vis_img = VisualizePointsClass(pts, output_img_h, output_img_w, d=d, theta=theta, vis_point_radius=vis_point_radius,
#                                Bbox=Bbox)
# SaveVisImg(vis_img, './1.png', './')

# alone = os.listdir(data_file)

# alone = os.listdir(data_file)
def LabelPoints2Images(data_file,scene_param,label_flag,LorR):
    pkg_list = alone  # os.listdir(data_file)
    for pkg_list_ in pkg_list:
        if LorR==0:
            pts_list_pr = os.listdir(data_file + '/' + pkg_list_ + '/pr_label_json')
        else:
            pts_list_pr = os.listdir(data_file + '/' + pkg_list_ + '/pr_radar_label_json')
        pts_list_gt = os.listdir(data_file + '/' + pkg_list_ + '/gt_label_json')
        pts_list_ld = os.listdir(data_file + '/' + pkg_list_ + '/lidar')
        pts_list = np.intersect1d(pts_list_pr, pts_list_gt)
        for i in range(pts_list.__len__()):
            pts_list[i]=pts_list[i][:-5]
        for i in range(pts_list_ld.__len__()):
            pts_list_ld[i]=pts_list_ld[i][:-4]
        pts_list = np.intersect1d(pts_list, pts_list_ld)
        pts_list=sorted(pts_list)
        count=0
        for pts in pts_list:
            count+=1
            pt_file = data_file + '/' + pkg_list_ + '/lidar/' + pts+'.bin'
            label_txt = data_file + '/' + pkg_list_ + '/label/' + pts + '.json'
            if LorR == 0:
                label_json = data_file + '/' + pkg_list_ + '/pr_label_json/' + pts + '.json'
            else:
                label_json = data_file + '/' + pkg_list_ + '/pr_radar_label_json/' + pts + '.json'
            gtlabel=data_file + '/' + pkg_list_ + '/gt_label_json/' + pts + '.json'
            pkl_label = data_file + '/' + pkg_list_ + '/pkl_label_json/' + pts + '.json'
            if label_flag==1:
                out_lidar_path = data_file + '/' + pkg_list_ + '/out_lidar_pre/'
            elif label_flag==2:
                out_lidar_path = data_file + '/' + pkg_list_ + '/out_lidar_gt/'
            elif label_flag==3:
                out_lidar_path = data_file + '/' + pkg_list_ + '/out_lidar_pkl/'
            else:
                if LorR == 0:
                    out_lidar_path = data_file + '/' + pkg_list_ + '/out_lidar/'
                else:
                    out_lidar_path = data_file + '/' + pkg_list_ + '/out_radar/'
            if not os.path.exists(out_lidar_path):
                os.makedirs(out_lidar_path)
            out_lidar_path_name=out_lidar_path + pts + '.jpg'
            label_file = label_json
            gtlabel_file = gtlabel
            start_time=time.time()
            # try:
            LabelPoints2ImagebyMaya(out_lidar_path_name=out_lidar_path_name, label_file=label_file,
                                            gtlabel=gtlabel_file, pt_file=pt_file,label_flag=label_flag,
                                            azimuth=scene_param["azimuth"], elevation=scene_param["elevation"],
                                            distance=scene_param["distance"],roll=scene_param["roll"],
                                            focalpoint=scene_param["focalpoint"],size=scene_param["size"])
            # except Exception as e:
            #     print(e)

            end_time = time.time()
            print('out_lidar_path_name',out_lidar_path_name,pts,':elapsed time is',end_time-start_time,'(s)')

def ToVideo(data_file,label_flag):
    pkg_list = alone  # ['data71_2020_05_13']#os.listdir(data_file)
    for pkg_list_ in pkg_list:
        if not os.path.exists(data_file + pkg_list_ + '/out_video/'):
            os.makedirs(data_file + pkg_list_ + '/out_video/')
        if label_flag==1:
            Pic2Video(data_file + pkg_list_ + '/out_lidar_pre/',
                      data_file + pkg_list_ + '/out_video/' + pkg_list_ + '_lidar_pre.mp4')
        elif label_flag==2:
            Pic2Video(data_file + pkg_list_ + '/out_lidar_gt/',
                      data_file + pkg_list_ + '/out_video/' + pkg_list_ + '_lidar_gt.mp4')
        elif label_flag==3:
            Pic2Video(data_file + pkg_list_ + '/out_lidar_pkl/',
                      data_file + pkg_list_ + '/out_video/' + pkg_list_ + '_lidar_pkl.mp4')
        else:
            Pic2Video(data_file + pkg_list_ + '/out_lidar/',
                      data_file + pkg_list_ + '/out_video/' + pkg_list_ + '_lidar.mp4')
        Pic2Video(data_file + pkg_list_ + '/image/', data_file + pkg_list_ + '/out_video/' + pkg_list_ + '_camera.mp4')

def resize_keep_aspectratio(image_src, dst_size):
    src_h, src_w = image_src.shape[:2]
    # print(src_h, src_w)
    dst_h, dst_w = dst_size

    # 判断应该按哪个边做等比缩放
    h = dst_w * (float(src_h) / src_w)  # 按照ｗ做等比缩放
    w = dst_h * (float(src_w) / src_h)  # 按照h做等比缩放

    h = int(h)
    w = int(w)

    if h <= dst_h:
        image_dst = cv2.resize(image_src, (dst_w, int(h)))
    else:
        image_dst = cv2.resize(image_src, (int(w), dst_h))

    h_, w_ = image_dst.shape[:2]
    # print(h_, w_)

    top = int((dst_h - h_) / 2);
    down = int((dst_h - h_ + 1) / 2);
    left = int((dst_w - w_) / 2);
    right = int((dst_w - w_ + 1) / 2);

    value = [0, 0, 0]
    borderType = cv2.BORDER_CONSTANT
    # print(top, down, left, right)
    image_dst = cv2.copyMakeBorder(image_dst, top, down, left, right, borderType, None, value)

    return image_dst

def video_fusion(data_file,image1,image2,LorR):
    fps = 10  # 视频每秒10帧
    size = (1600, 970)  # 需要转为视频的图片的尺寸
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    pkg_list = alone  # os.listdir(data_file)
    for pkg_list_ in pkg_list:
        pts_list = os.listdir(data_file + '/' + pkg_list_ + image1)
        if LorR==0:
            video_name = data_file+ '/' + pkg_list_ + '/out_video/' + pkg_list_ + '_lidar_fusion.mp4'
        else:
            video_name = data_file + '/' + pkg_list_ + '/out_video/' + pkg_list_ + '_radar__fusion.mp4'
        if not os.path.exists(data_file + pkg_list_ + '/out_video/'):
            os.makedirs(data_file + pkg_list_ + '/out_video/')
        pts_list = sorted(pts_list)
        video_name2 = data_file + '/' + pkg_list_ + image1 + pts_list[0]
        frame2 = cv2.imread(video_name2)
        size = (frame2.shape[1], frame2.shape[0])
        videoWriter = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, size)
        for pts in pts_list:
            video_name2 = data_file + '/' + pkg_list_ + image1 + pts
            video_name1 = data_file + '/' + pkg_list_ + image2 + pts
            print(video_name1,video_name2)
            frame2 = cv2.imread(video_name2)
            frame1 = cv2.imread(video_name1)
            # frame11 = frame1[128:-128, :, :]
            frame1_temp = resize_keep_aspectratio(frame1, (500, 700))
            frame2[0:500, 1860:2560, :] = frame1_temp
            videoWriter.write(frame2)
        videoWriter.release()
        print('video_name:',video_name)

# alone = ['data_074']
# alone = ['data36_2020_05_12']
# alone = ['data15_2020_05_11']
# alone = ['data57_2020_05_12']
# alone = ['data76_2020_05_13']
# alone = ['data95_2020_05_14']
# alone = ['040302_release_rain_truck']
# data_file = './dataset/'
# label_flag=4
# LabelPoints2Images(data_file,scene_param,label_flag)
# video_fusion(data_file)
# ToVideo(data_file,label_flag)

# file='/media/king/MyPassport0/livox_truck_dataset_v2.0/data/'
# alones=os.listdir(file)
# alone=[]
# for alone_ in alones:
#     if '.' in alone_:
#         continue
#     alone.append(alone_)
# alone=['data9_2021_03_29']
# id=['11','12','14','15','16','17','18','19','20','21','22','23','24','25','26','27','28']
# alone=[]
# for i in range(id.__len__()):
#     alone.append('data'+id[i]+'_2021_04_03')
# alone=alone[:17]
# data_file=file

data_file='/home/king/workplace/data'
# alone=['data14_2021_04_03','data27_2021_04_03','data28_2021_04_03']
alone=['data27_2021_04_03']
label_flag=4
LorR=1
LabelPoints2Images(data_file,scene_param,label_flag,LorR)
# video_fusion(data_file,'/out_lidar/','/image/',LorR)
# alone=['data27_2021_04_03','data28_2021_04_03']
# label_flag=4
# LorR=1
# # LabelPoints2Images(data_file,scene_param,label_flag,LorR)
# video_fusion(data_file,'/out_radar/','/image/',LorR)